"""
指数平滑 & 预测 & 评估
"""
import numpy as np
import pandas as pd
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import TimeSeriesSplit
from statsmodels.tsa.api import SimpleExpSmoothing


def exp_sm(series, alpha):
    """
    指数平滑
    :param series:
    :param alpha:
    :return:
    """
    # 一次平滑
    rs_1 = [series[0]]
    # 二次平滑
    rs_2 = [alpha * rs_1[-1]]
    # 三次平滑
    rs_3 = [alpha * rs_2[-1]]

    # 二次平滑参数
    lr2_intercept = [2 * rs_1[-1] - rs_2[-1]]
    lr2_weight = [(rs_1[-1] - rs_2[-1]) * alpha / (1 - alpha)]

    # 三次平滑参数
    lr3_intercept = [3 * rs_1[-1] - 3 * rs_2[-1] + rs_3[-1]]
    lr3_weight = \
        [((6 - 5 * alpha) * rs_1[-1] - (10 - 8 * alpha) * rs_2[-1] +
          (4 - 3 * alpha) * rs_3[-1]) *
         alpha / (2 * np.power(1 - alpha, 2))]
    lr3_square_weight = \
        [(rs_1[-1] - 2 * rs_2[-1] + rs_3[-1]) *
         np.power(alpha, 2) / np.power(1 - alpha, 2)]

    for idx_num in range(1, len(series)):
        # 下一期的平滑值
        rs_1.append(alpha * series[idx_num - 1] +
                    (1 - alpha) * rs_1[-1])
        rs_2.append(alpha * rs_1[-1] + (1 - alpha) * rs_2[-1])
        rs_3.append(alpha * rs_2[-1] + (1 - alpha) * rs_3[-1])

        lr2_intercept.append(2 * rs_1[-1] - rs_2[-1])
        lr2_weight.append((rs_1[-1] - rs_2[-1]) * alpha / (1 - alpha))

        lr3_intercept.append(3 * rs_1[-1] - 3 * rs_2[-1] + rs_3[-1])
        lr3_weight.append(
            ((6 - 5 * alpha) * rs_1[-1] - (10 - 8 * alpha) * rs_2[-1] +
             (4 - 3 * alpha) * rs_3[-1]) *
            alpha / (2 * np.power(1 - alpha, 2)))
        lr3_square_weight.append(
            (rs_1[-1] - 2 * rs_2[-1] + rs_3[-1]) *
            np.power(alpha, 2) / np.power(1 - alpha, 2))

    return [rs_1, rs_2, rs_3, lr2_intercept, lr2_weight,
            lr3_intercept, lr3_weight, lr3_square_weight]


def predict(series, alpha, period):
    """
    预测
    :param series:
    :param alpha:
    :param period: 预测超前期数
    :return:
    """
    _, _, _, lr2_a, lr2_b, lr3_a, lr3_b, lr3_c = exp_sm(series=series,
                                                        alpha=alpha)
    # 二次平滑预测
    pred_exp_2 = lr2_a[-1] + lr2_b[-1] * period
    # 三次平滑预测
    pred_exp_3 = \
        lr3_a[-1] + lr3_b[-1] * period + \
        0.5 * lr3_c[-1] * np.power(period, 2)
    return pred_exp_2, pred_exp_3


def evalution_mape(y, pred):
    """
    MAPE评估
    :param y:
    :param pred:
    :return:
    """
    mape_sum = 0
    length = len(y)
    for idx, num in enumerate(y):
        mape_sum += np.abs((num - pred[idx]) / num) / length
    return mape_sum


def exp_sm_main(series, alpha):
    """
    主函数
    :param series:
    :param alpha: 超参
    :return:
    """
    if isinstance(series, pd.Series):
        data_set = series.to_numpy()

    if not isinstance(series, np.ndarray):
        data_set = np.array(series)
    else:
        data_set = series

    # rms评估
    rms_simple = []
    rms_loss_2 = []
    rms_loss_3 = []
    # mape评估
    mape_simple = []
    mape_loss_2 = []
    mape_loss_3 = []

    tscv = TimeSeriesSplit()
    for train_idx, test_idx in tscv.split(data_set):
        train, test = data_set[train_idx], data_set[test_idx]
        peroid = len(test_idx)
        y2_pred = []
        y3_pred = []

        # 简单指数平滑
        simple_model = SimpleExpSmoothing(data_set).fit(smoothing_level=alpha)
        y_pred = simple_model.forecast(peroid)
        rms_simple.append(np.sqrt(mean_squared_error(test, y_pred)))
        mape_simple.append(evalution_mape(test, y_pred))

        # 三种平滑
        sm1, sm2, sm3, sm2_a, sm2_b, sm3_a, sm3_b, sm3_c = \
            exp_sm(series=train, alpha=alpha)
        # predict
        for num in range(peroid):
            pred_2, pred_3 = predict(series=train, alpha=alpha, period=num)
            sm2.append(pred_2)
            sm3.append(pred_3)

            y2_pred.append(pred_2)
            y3_pred.append(pred_3)

        # 评估
        rms_loss_2.append(np.sqrt(mean_squared_error(test, y2_pred)))
        mape_loss_2.append(evalution_mape(test, y2_pred))
        rms_loss_3.append(np.sqrt(mean_squared_error(test, y3_pred)))
        mape_loss_3.append(evalution_mape(test, y3_pred))

    print('rms: exp_2: {}, exp_3: {}, simple: {}'.format(np.mean(rms_loss_2),
                                                         np.mean(rms_loss_3),
                                                         np.mean(rms_simple)))
    print('mape: exp_2: {}, exp_3: {}, simple: {}'.format(np.mean(mape_loss_2),
                                                          np.mean(mape_loss_3),
                                                          np.mean(mape_simple)))
